import json
from typing import List, Dict


def load_strategy_from_file(file_path: str) -> List[Dict]:
    """从JSON文件加载复杂并行策略配置"""
    with open(file_path, 'r') as f:
        config = json.load(f)

    # 转换格式适配计算函数
    stages = []
    for stage in config['stages']:
        stages.append({
            'start_layer': stage['layer_range'][0],
            'end_layer': stage['layer_range'][1],
            'submesh_shape': tuple(stage['submesh_shape']),
            'strategies': stage['strategies']
        })
    return stages


def calculate_manual_pipeline_cost(
    L: int,                   # 模型总层数（例如：12）
    S: int,                    # 流水线并行度数（PP，例如：2）
    d: int,                    # 数据并行度数（DP，例如：2）
    t: int,                    # 张量并行度数（TP，例如：2）
    num_devices: int,          # 设备总数（需满足 S*d*t = num_devices，例如：8）
    num_microbatches: int,     # 微批次数量（例如：32）
    device_tflops: float,      # 单设备算力（单位：TFLOPS，例如：100）
    bandwidth_gbps: float,     # 设备间带宽（单位：GB/s，例如：200）
    layer_flops: float,        # 单层计算量（单位：FLOPS，例如：1e12）
    parameter_size: float,     # 模型参数量（单位：字节，例如：1e9）
    gradient_size: float       # 梯度大小（单位：字节，例如：2e9）
) -> float:
    """
    计算人工配置的流水线并行策略的总时间开销
    """
    # 1. 检查设备数量是否合法
    if S * d * t != num_devices:
        raise ValueError(f"S*d*t ({S}*{d}*{t}) != num_devices ({num_devices})")

    # 2. 均匀划分模型层到每个stage
    layers_per_stage = L // S

    # 3. 计算单个stage的计算时间（前向+反向）
    # 转换TFLOPS到FLOPS（1 TFLOPS = 1e12 FLOPS）
    device_flops = device_tflops * 1e12
    forward_time = (layer_flops * layers_per_stage) / device_flops  # 单位：秒
    backward_time = 2 * forward_time  # 反向传播时间通常是前向的2倍
    compute_time = forward_time + backward_time

    # 4. 计算通信时间（TP和DP）
    # 转换带宽单位（1 GB/s = 1e9 字节/秒）
    bandwidth = bandwidth_gbps * 1e9
    # 张量并行通信时间（参数同步）
    tp_comm_time = (parameter_size / t) / bandwidth
    # 数据并行通信时间（梯度同步）
    dp_comm_time = (gradient_size / d) / bandwidth
    total_comm_time = tp_comm_time + dp_comm_time

    # 5. 单个stage的总时间（假设所有stage均匀分配）
    stage_time = compute_time + total_comm_time

    # 6. 总时间 = 首阶段时间 + (微批次数-1) * 最慢阶段时间
    # 因人工等分策略，所有stage时间相同，故直接使用stage_time
    total_cost = stage_time + (num_microbatches - 1) * stage_time

    return total_cost


def calculate_complex_pipeline_cost(
        stages_config: List[Dict],  # 各阶段配置信息
        num_microbatches: int,  # 微批次数量
        device_tflops: float,  # 单设备算力（TFLOPS）
        bandwidth_gbps: float,  # 带宽（GB/s）
        layer_flops: List[float],  # 每层计算量（FLOPS）
        parameter_size: List[float],  # 每层参数量（字节）
        gradient_size: List[float],  # 每层梯度大小（字节）
) -> float:
    """
    计算复杂并行策略的总时间开销
    输入参数说明：
    stages_config: [
        {
            'start_layer': 0,             # 阶段起始层（包含）
            'end_layer': 2,               # 阶段结束层（不包含）
            'submesh_shape': (4, 8),      # 设备子网格形状
            'strategies': [                # 每层的并行策略（按阶段内顺序）
                {'dp': 32, 'tp': 1},      # 第0层策略
                {'dp': 8, 'tp': 1},       # 第1层策略
            ]
        },
        # ...其他阶段配置
    ]
    """
    stage_times = []
    for stage in stages_config:
        start = stage['start_layer']
        end = stage['end_layer']
        submesh_shape = stage['submesh_shape']
        strategies = stage['strategies']
        num_layers = end - start
        # assert len(strategies) == num_layers, f"阶段{start}-{end}策略数量与层数不符"

        # 验证子网格设备总数
        submesh_devices = submesh_shape[0] * submesh_shape[1]

        # 计算阶段的计算时间
        compute_time = 0.0
        for i in range(num_layers):
            layer_idx = start + i
            dp = strategies[i]['dp']
            tp = strategies[i]['tp']

            # 检查DP*TP是否等于子网格设备数
            assert dp * tp == submesh_devices, (
                f"阶段{start}-{end}层{i}：DP({dp})*TP({tp}) != "
                f"子网格设备数{submesh_devices}"
            )

            # 计算该层的前向+反向时间
            device_flops = device_tflops * 1e12  # 转换为FLOPS
            layer_compute = layer_flops[layer_idx] / tp  # 张量并行分摊计算量
            forward_time = layer_compute / device_flops
            backward_time = 2 * forward_time  # 反向时间为前向2倍
            compute_time += forward_time + backward_time

        # 计算阶段的通信时间
        comm_time = 0.0
        for i in range(num_layers):
            layer_idx = start + i
            dp = strategies[i]['dp']
            tp = strategies[i]['tp']

            # 张量并行通信（参数同步）
            tp_comm = (parameter_size[layer_idx] / tp) / (bandwidth_gbps * 1e9)
            # 数据并行通信（梯度同步）
            dp_comm = (gradient_size[layer_idx] / dp) / (bandwidth_gbps * 1e9)
            comm_time += tp_comm + dp_comm

        # 阶段总时间（计算+通信）
        stage_time = compute_time + comm_time
        stage_times.append(stage_time)

    # 流水线总时间 = 首尾耗时 + (微批次数-1)*最慢阶段耗时
    total_time = sum(stage_times) + (num_microbatches - 1) * max(stage_times)
    return total_time


# 测试案例（基于论文典型参数）
if __name__ == "__main__":
    file_path = '/home/lthpc/nvmessd/wangjiaqian/simulator_by_alpa_layer/experiment/complex_strategy.json'

    # 全局参数（与简单策略保持一致）
    L = 24  # 模型总层数（示例配置）
    num_microbatches = 32
    device_tflops = 100  # 100 TFLOPS/GPU
    bandwidth_gbps = 200  # 200 GB/s
    layer_flops = [1e12] * L  # 每层计算量（简化假设）
    parameter_size = [1e9] * L  # 每层参数量（1GB）
    gradient_size = [2e9] * L  # 每层梯度大小（2GB）

    # 从文件加载复杂策略
    complex_stages = load_strategy_from_file(file_path)
    complex_time = calculate_complex_pipeline_cost(
        complex_stages, num_microbatches,
        device_tflops, bandwidth_gbps,
        layer_flops, parameter_size, gradient_size
    )

    simple_time = calculate_manual_pipeline_cost(
        L=24, S=4, d=2, t=4,  # 假设均匀划分
        num_devices=4 * 2 * 4,
        num_microbatches=num_microbatches,
        device_tflops=device_tflops,
        bandwidth_gbps=bandwidth_gbps,
        layer_flops=1e12,
        parameter_size=1e9,
        gradient_size=2e9
    )

    # 结果分析
    print("\n===== 并行策略对比分析 =====")
    print(f"[复杂策略] 总时间: {complex_time:.2f}秒")
    print(f"[简单策略] 总时间: {simple_time:.2f}秒")
    print(f"时间减少: {simple_time - complex_time:.2f}秒 ({((simple_time - complex_time) / simple_time * 100):.1f}%)")

    # 各阶段耗时分析
    print("\n[复杂策略各阶段耗时]")
    for i, stage in enumerate(complex_stages):
        stage_time = calculate_complex_pipeline_cost(
            [stage], 1,  # 计算单个阶段耗时
            device_tflops, bandwidth_gbps,
            layer_flops, parameter_size, gradient_size
        )
        print(f"Stage {i}: 层{stage['start_layer']}-{stage['end_layer']} | "
              f"子网格{stage['submesh_shape']} | 耗时: {stage_time:.3f}秒")

